package tcpserver

import (
	"context"
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

var (
	ErrServerClosed = errors.New("tcp: Server Closed")
	ErrAbortHandler = errors.New("tcp: abort TCPHandler")
	ServerContextKey = &contextKey{"tcp-server"}
	LocalAddrContextKey = &contextKey{"local-addr"}
)

type onceCloseListener struct{
	net.Listener
	once sync.Once
	closeErr error
}

func (oc *onceCloseListener) Close() error{
	oc.once.Do(oc.close)
	return oc.closeErr
}

func (oc *onceCloseListener) close(){
	oc.closeErr = oc.Listener.Close()
}

type TCPHandler interface{
	ServeTCP(ctx context.Context, conn net.Conn)
}

type TCPServer struct {
	Addr string
	Handler TCPHandler
	err error
	BaseCTX context.Context

	WriteTimeout time.Duration
	ReadTimeout time.Duration
	KeepAliveTimeout time.Duration

	mu sync.Mutex
	inShutdown int32
	doneChan chan struct{}
	l *onceCloseListener
}

func (s *TCPServer) shutDown()bool{
	return atomic.LoadInt32((&s.inShutdown)) != 0
}

func (srv *TCPServer) ListenAndServe()error{
	if srv.shutDown(){
		return ErrServerClosed
	}
	if srv.doneChan == nil{
		srv.doneChan = make(chan struct{})
	}
	addr := srv.Addr
	if addr == "" {
		return errors.New("need addr")
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil{
		return err
	}
	return srv.Serve(tcpKeepAliveListener{ln.(*net.TCPListener)})
}
func ListenAndServe(addr string, handler TCPHandler) error {
	server := &TCPServer{Addr: addr, Handler: handler, doneChan: make(chan struct{})}
	return server.ListenAndServe()
}

func (s *TCPServer) getDoneChan() <- chan struct{}{
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.doneChan == nil{
		s.doneChan = make(chan struct{})
	}
	return s.doneChan
}

func (srv *TCPServer) newConn(rwc net.Conn) *conn {
	c := &conn{
		server: srv,
		rwc: rwc,
	}
	if d := c.server.ReadTimeout; d != 0 {
		c.rwc.SetReadDeadline(time.Now().Add(d))
	}
	if d := c.server.WriteTimeout; d != 0 {
		c.rwc.SetWriteDeadline(time.Now().Add(d))
	}
	if d := c.server.KeepAliveTimeout; d != 0 {
		if tcpConn, ok := c.rwc.(*net.TCPConn); ok {
			tcpConn.SetKeepAlive(true)
			tcpConn.SetKeepAlivePeriod(d)
		}
	}
	return c
}

func (srv *TCPServer) Serve(l net.Listener) error{
	srv.l = &onceCloseListener{Listener: l}
	defer srv.l.Close()
	if srv.BaseCTX == nil{
		srv.BaseCTX = context.Background()
	}
	baseCTX := srv.BaseCTX
	ctx := context.WithValue(baseCTX, ServerContextKey, srv)
	for {
		rw, e := l.Accept()
		if e != nil{
			select{
			case <- srv.getDoneChan():
				return ErrServerClosed
			default:
			}
			fmt.Printf("accept fail, err: %v\n", e)
		}
		c := srv.newConn(rw)
		go c.serve(ctx)
	}

}

func (srv *TCPServer) Close() error{
	atomic.StoreInt32(&srv.inShutdown,1)
	close(srv.doneChan)
	srv.l.Close()
	return nil
}